import glob
import shutil
from xml.etree import cElementTree
import numpy as np
import os


class DataSpliter:
    def __init__(self, data_dir, class_list, split_ratio_dict, random_str):
        self.data_dir = data_dir
        self.class_list = class_list
        self.split_ratio_dict = split_ratio_dict
        self.random_str = random_str

    def get_random(self):
        cnt = 0
        ratio = self.__split_ratio()
        for num in ratio:
            cnt += int(num*num)
        for i in self.random_str:
            cnt += ord(i)
        return cnt

    def __split_ratio(self):
        return list(self.split_ratio_dict.values())

    def data_split(self):
        class_dataset = []
        xml_names = glob.glob(self.data_dir + '/*.xml')
        for xml_name in range(len(xml_names)):
            img_path = xml_names[xml_name].split('.')[0] + '.jpg'
            root_element = cElementTree.parse(xml_names[xml_name])
            root = root_element.getroot()
            for obj in root.iter('object'):
                box_name = obj.find('name')
                key = str(box_name.text)
                if key in self.class_list:
                    if [img_path, xml_names[xml_name]] not in class_dataset:
                        if {'image': img_path, 'xml': xml_names[xml_name]} not in class_dataset:
                            class_dataset.append({'image': img_path, 'xml': xml_names[xml_name]})
        np.random.seed(self.get_random())
        np.random.shuffle(class_dataset)
        split_ratio = self.__split_ratio()
        split_name = list(self.split_ratio_dict.keys())
        datas = {}
        flag = 0
        for i in range(len(split_ratio)):
            if i == 0:
                flag += int(split_ratio[i] / sum(split_ratio)*len(class_dataset))
                data = class_dataset[:flag]
                datas[split_name[i]] = data
            elif i == len(split_ratio)-1:
                data = class_dataset[flag:len(class_dataset)]
                datas[split_name[i]] = data
            else:
                data = class_dataset[flag:int(split_ratio[i] / sum(split_ratio)*len(class_dataset)) + flag]
                flag = int(split_ratio[i] / sum(split_ratio)*len(class_dataset)) + flag
                datas[split_name[i]] = data
        return datas


def save(save_path_dict, data_split_dict):
    for key, value in data_split_dict.items():
        for key_0, value_0 in save_path_dict.items():
            if key == key_0:
                for old_path in value:
                    new_img_path = value_0 + '\\' + old_path['image'].split('\\')[-1]
                    new_xml_path = value_0 + '\\' + old_path['xml'].split('\\')[-1]
                    if not os.path.exists(value_0):
                        os.makedirs(value_0)
                    shutil.copy(old_path['image'], new_img_path)
                    shutil.copy(old_path['xml'], new_xml_path)


def data_split(data_dir, class_list, split_ratio_dict, random_str):
    data_sp = DataSpliter(data_dir, class_list, split_ratio_dict, random_str)
    return data_sp.data_split()
